# -*- coding: utf-8 -*-
# @Time : 2021/11/18 13:49
# @Author : Xiang Qian Xiang Qian
# @Email : qianxljp@126.com
# @File : hrrp_preprocessor.py
# @Project : hrrp_jt

import pandas as pd
import scipy.io as sco
import os

IMB = {
    # imb_level:[ball,cabin,simple,high,warhead]
    'Im0': [1.0, 1.0, 1.0, 1.0, 1.0],
    'Im1': [1.0, 0.95, 0.9, 0.85, 0.8],
    'Im2': [1.0, 0.9, 0.8, 0.7, 0.6],
    'Im3': [1.0, 0.9, 0.8, 0.6, 0.4],
    'Im4': [1.0, 0.6, 0.5, 0.4, 0.3],

}

class Hrrp_preprocessor():

    def __init__(self, src=r'/soar/data/hrrp_0.05/', dst=r'/soar/data/hrrp_0.05/preprocessed/', snr='snr0',
                 imb='Im0', ):
        self.snr = snr
        self.imb = imb
        self.src = src
        self.dst = dst
        self.random_state = 2021
        self.preprocessing()

    def preprocessing(self):
        self.data_id = '{}_{}'.format(self.imb, self.snr)
        self.hrrp_mat2pd()
        self.concat_targets()

    def hrrp_mat2pd(self):
        snr = self.snr
        data_path = self.src
        if snr == 'non':
            self.ball = pd.DataFrame(sco.loadmat(data_path + "ball_Decoy_HH_0_05_180.mat")['ball_Decoy_HH_0_05_180'])
            self.cabin = pd.DataFrame(
                    sco.loadmat(data_path + "mother_Cabin_HH_0_05_180.mat")['mother_Cabin_HH_0_05_180'])
            self.simple = pd.DataFrame(
                    sco.loadmat(data_path + "simple_Decoy_HH_0_05_180.mat")['simple_Decoy_HH_0_05_180'])
            self.high = pd.DataFrame(
                    sco.loadmat(data_path + "hight_imitation_Decoy_HH_0_05_180.mat")[
                        'hight_imitation_Decoy_HH_0_05_180'])
            self.warhead = pd.DataFrame(sco.loadmat(data_path + "warhead_HH_0_05_180.mat")['warhead_HH_0_05_180'])
        # 加载噪声数据
        else:
            data_path = data_path + 'hrrp_noised/'
            self.ball = pd.DataFrame(sco.loadmat(data_path + "{}/{}_b.mat".format(snr, snr))['b'])
            self.high = pd.DataFrame(sco.loadmat(data_path + "{}/{}_h.mat".format(snr, snr))['h'])
            self.cabin = pd.DataFrame(sco.loadmat(data_path + "{}/{}_m.mat".format(snr, snr))['m'])
            self.simple = pd.DataFrame(sco.loadmat(data_path + "{}/{}_s.mat".format(snr, snr))['s'])
            self.warhead = pd.DataFrame(sco.loadmat(data_path + "{}/{}_w.mat".format(snr, snr))['w'])

        # self.z_score()
        # self.min_max()
        self.ball.insert(loc=256, column='label', value=0)
        self.cabin.insert(loc=256, column='label', value=1)
        self.simple.insert(loc=256, column='label', value=2)
        self.high.insert(loc=256, column='label', value=3)
        self.warhead.insert(loc=256, column='label', value=4)
        assert self.ball.shape == self.cabin.shape == self.simple.shape == self.high.shape == self.warhead.shape
        # print(self.ball.shape)

    def z_score(self):
        all = pd.concat([self.ball, self.cabin, self.simple, self.high, self.warhead])
        mean = all.mean()
        std = all.std()
        self.ball = (self.ball - mean) / std
        self.cabin = (self.cabin - mean) / std
        self.simple = (self.simple - mean) / std
        self.high = (self.high - mean) / std
        self.warhead = (self.warhead - mean) / std

    def min_max(self):
        all = pd.concat([self.ball, self.cabin, self.simple, self.high, self.warhead])
        min = all.min()
        max = all.max()
        self.ball = (self.ball - min) / (max - min)
        self.cabin = (self.cabin - min) / (max - min)
        self.simple = (self.simple - min) / (max - min)
        self.high = (self.high - min) / (max - min)
        self.warhead = (self.warhead - min) / (max - min)

    def sample_target(self, target, sample_rate=1):
        # 先打乱数据
        # print(len(target))
        target = target.sample(frac=1, random_state=self.random_state, replace=True)
        target.reset_index(drop=True, inplace=True)
        # 80%的数据作为训练集
        target_train_all = target.sample(frac=0.8, random_state=self.random_state)
        target_test = target.drop(index=target_train_all.index)

        # 抽取训练集中的数据得到不平衡数据集
        target_train = target_train_all.sample(frac=sample_rate, random_state=self.random_state)
        dataset = {
            'train': target_train,
            'test': target_test,
            'whole': pd.concat([target_train, target_test]),
        }
        return dataset

    def concat_targets(self):
        # rates用于构造不平衡数据集
        rates = IMB[self.imb]
        d_ball = self.sample_target(self.ball, sample_rate=rates[0])
        d_cabin = self.sample_target(self.cabin, sample_rate=rates[1])
        d_simple = self.sample_target(self.simple, sample_rate=rates[2])
        d_high = self.sample_target(self.high, sample_rate=rates[3])
        d_warhead = self.sample_target(self.warhead, sample_rate=rates[4])
        for d_name in ['train', 'test', 'whole']:
            ball = d_ball[d_name]
            cabin = d_cabin[d_name]
            simple = d_simple[d_name]
            high = d_high[d_name]
            warhead = d_warhead[d_name]
            if self.imb == 'Im0':
                assert ball.shape == cabin.shape == simple.shape == high.shape == warhead.shape
                print('Shape of each target in {} dataset: {}'.format(d_name, ball.shape))
            dset = pd.concat([ball, cabin, simple, high, warhead])
            # print(dset['label'].value_counts())
            dset.reset_index(drop=True, inplace=True)
            # print(dset['label'].value_counts())
            dset = dset.sample(frac=1.0, random_state=self.random_state)
            # print(dset['label'].value_counts())
            dset.reset_index(drop=True, inplace=True)
            path = self.dst + self.data_id + '/'
            if not os.path.exists(path):
                os.makedirs(path)
                print('Make path: {}'.format(path))
            dset.to_pickle(path + self.data_id + '_' + d_name + '.pkl')
            print('Shape of {} dataset: {}'.format(d_name, dset.shape))
            print(dset['label'].value_counts())
        return dset

for snr in ["non",
            "snr-10", "snr-5", "snr0", "snr5", "snr10", "snr15", "snr20", "snr25", "snr30", "snr35", "snr40"
            ]:
    print('\n' + 30 * '=' + snr + 30 * '=')
    for imb in ['Im0',
                'Im1', 'Im2', 'Im3', 'Im4'
                ]:
        a = Hrrp_preprocessor(snr=snr, imb=imb)
